package com.zl.Util;

import com.jcraft.jsch.*;

import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Map;

import static java.lang.System.in;

/**
 * @author zhoul
 * @since 2016-09-07-10:33
 */
public class SSHCmdUtil {

    private static Map<String, Session> pool = new HashMap<>();

    public static Session getSession(String userName, String passwd, String addr, int port) throws JSchException {
        Session session;
        if (pool.containsKey(addr) && pool.get(addr).isConnected()) {
            session = pool.get(addr);
        } else {
            JSch jSch = new JSch();
            session = jSch.getSession(userName, addr, port);
            session.setPassword(passwd);
            session.setConfig("StrictHostKeyChecking", "no");
            pool.put(addr, session);
            session.connect(1000);
        }
        return session;
    }
    private static String read(InputStream in, Channel channel) throws IOException {
        StringBuilder result = new StringBuilder();
        byte[] tmp = new byte[1024];
        while (true) {
            while (in.available() > 0) {
                int i = in.read(tmp, 0, 1024);
                if (i < 0) break;
                result.append(new String(tmp, 0, i, "utf-8"));
            }
            if (channel.isClosed()) {
                if (in.available() > 0) continue;
//                    System.out.println("exit-status: " + channel.getExitStatus());
//                    resCode = channel.getExitStatus();
                break;
            }
        }
        return result.toString();
    }

    //使用nohup时，cmd的形式需要是：  nohup java -jar XXX.jar > XXX.log 2>&1 &
    public static String runCmdWithNohup(String userName, String passwd, String addr, int port, String cmd) throws IOException, JSchException {
        String result;
        Channel channel = null;
        Session session;
        try {
            session = getSession(userName, passwd, addr, port);

            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(cmd);
            ((ChannelExec) channel).setPty(false);
            channel.setInputStream(null);
            ((ChannelExec) channel).setErrStream(System.err);
            InputStream in = channel.getInputStream();
            channel.connect();

            result = read(in, channel);
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
        }

        return result;
    }

    //执行简单的cmd命令：适用于“立即返回结果”的命令。不关闭会话
    public static String runCmdWithSession(String userName, String passwd, String addr, int port, String cmd) throws JSchException, IOException {
        String result;
        Channel channel = null;
        Session session;
        try {
            session = getSession(userName, passwd, addr, port);

            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(cmd);
//            ((ChannelExec) channel).setPty(false);
            channel.setInputStream(null);
            ((ChannelExec) channel).setErrStream(System.err);
            InputStream in = channel.getInputStream();
            channel.connect();

            result = read(in, channel);
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
        }
        return result;
    }

    //执行简单的cmd命令：适用于“立即返回结果”的命令 。关闭会话
    public static String runCmdWithNoSession(String userName, String passwd, String addr, int port, String cmd) throws JSchException, IOException {
        String result;
        Channel channel = null;
        Session session = null;
        try {
            JSch jSch = new JSch();
            session = jSch.getSession(userName, addr, port);
            session.setPassword(passwd);
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();

            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(cmd);
            // channel.setXForwarding(true);
            //channel.setInputStream(System.in);
            channel.setInputStream(null);
            //channel.setOutputStream(System.out);
            //FileOutputStream fos=new FileOutputStream("/tmp/stderr");
            //((ChannelExec)channel).setErrStream(fos);
            ((ChannelExec) channel).setErrStream(System.err);
            InputStream in = channel.getInputStream();
            channel.connect();

            result = read(in, channel);

        } finally {
            if (channel != null) {
                channel.disconnect();
            }
            if (session != null) {
                session.disconnect();
            }
        }
        return result;
    }


}
